/*
 * PropagateImage.java
 *
 * Created on Jan 19, 2012
 *
 * Copyright (c) 2011 Artur Rataj.
 *
 * This code is distributed under the terms of the GNU Library
 * General Public License, either version 3 of the license or, at
 * your option, any later version.
 */

package pl.gliwice.iitis.spatialnn.propagate.image;

import java.awt.AlphaComposite;
import java.awt.Graphics2D;
import java.awt.RenderingHints;
import java.util.*;
import java.awt.image.BufferedImage;
import java.awt.image.Raster;
import java.awt.image.WritableRaster;
import java.io.File;
import java.io.IOException;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.imageio.ImageIO;

import pl.gliwice.iitis.spatialnn.network.*;
import pl.gliwice.iitis.spatialnn.propagate.*;

/**
 * Propagates images though a network, generates an output image.
 * 
 * @author Artur Rataj
 */
public class PropagateImage {
    /**
     * Type of images that can be propagated.
     */
    final static int IMAGE_TYPE = BufferedImage.TYPE_USHORT_GRAY;
    
    /**
     * Network, through which propagate the images.
     */
    Network network;
    /**
     * Propagation.
     */
    Propagate propagate;
    
    /**
     * Creates a new instance of <code>PropagateImage</code>.
     * 
     * @param network network to propagate through
     */
    public PropagateImage(Network network) {
        this.network = network;
        propagate = new Propagate(this.network);
    }
    /**
     * Returns a cache from the <code>Propagate</code> object.
     * 
     * @return a propagation cache
     */
    public Cache getCache() {
        return propagate.getCache();
    }
    /**
     * Scales the width of the input images to match resolutions of respective
     * input membranes. The height of all these images should still be equal, but it is
     * converted as well to maintain aspect ratio. This method also
     * converts the images to <code>BufferedImage.TYPE_USHORT_GRAY</code>
     * if necessary.
     * 
     * @param images images to scale and convert, keyed with keys of input membranes;
     * if an image would need conversion, it is not modified -- instead, a new images is
     * put into the map
     * @param heightScale how to  scale the height, that is, the number of samples;
     * typically &lt;= 1.0
     * @return old images not converted or new images converted; apt for <code>get()</code>
     */
    public void prepareImages(Map<String, BufferedImage> images,
            double heightScale) {
        int samplesNum = -1;
        for(String key : images.keySet()) {
            BufferedImage image = images.get(key);
            Membrane m = network.membraneMap.get(key);
            int width = m.getRes();
            boolean convert = m.type != MembraneType.IN ||
                    width != image.getWidth() ||
                    image.getType() != IMAGE_TYPE;
            if(samplesNum == -1)
                samplesNum = image.getHeight();
            else if(samplesNum != image.getHeight())
                throw new RuntimeException("image " + key + ": mismatch of the number of samples");
            if(convert) {
                int height = (int)Math.round(samplesNum*heightScale);
                BufferedImage resizedImage = resize(image, width, height);
                images.put(key, resizedImage);
                try {
                    ImageIO.write(resizedImage, "PNG", new File("/tmp/image" + key + ".png"));
                } catch (IOException e) {
                    throw new RuntimeException("io error : " + e.toString());
                }
            }
        }
    }
    /**
     * Propagates given signals through the network. All images should be
     * grayscale, of the type <code>BufferedImage.TYPE_USHORT_GRAY</code>.
     * 
     * @param in input images, keyed with membrane keys; must be of the width
     * being membrane resolutions, and of an equal height being the number of
     * samples
     * @param outScale scale of the grayscale samples; <code>1.0/outScale</code>
     * has a maximum brightness; the samples are clipped
     * @param cache cache to be used by <code>Propagate.get()</code>; see the docs
     * of that method for details
     * @return output image, of the width being the resolution of the output membrane;
     * and of the height being equal to the height of the input images
     */
    public BufferedImage get(Map<String, BufferedImage> in, double outScale, Cache cache) {
        int samplesNum = -1;
        Map<String, double[]> inBuffer = new HashMap<>();
        for(String key : in.keySet()) {
            BufferedImage image = in.get(key);
            Membrane m = network.membraneMap.get(key);
            if(m.type != MembraneType.IN)
                throw new RuntimeException("input membrane expected");
            if(m.getRes() != image.getWidth())
                throw new RuntimeException("image " + key + ": mismatch of the resolution");
            inBuffer.put(key, new double[image.getWidth()]);
            if(samplesNum == -1)
                samplesNum = image.getHeight();
            else if(samplesNum != image.getHeight())
                throw new RuntimeException("image " + key + ": mismatch of the number of samples");
            if(image.getType() != IMAGE_TYPE)
                throw new RuntimeException("image " + key + ": unsupported raster format");
        }
        BufferedImage out = new BufferedImage(cache.outM.getRes(), samplesNum,
                IMAGE_TYPE);
        final int MAX_SAMPLE = (1 << 16) - 1;
        for(int y = 0; y < samplesNum; ++y) {
            for(String key : in.keySet()) {
                BufferedImage image = in.get(key);
                Membrane m = network.membraneMap.get(key);
                int res = m.getRes();
                double[] buffer = inBuffer.get(key);
                Raster r = image.getRaster();
                for(int x = 0; x < res; ++x)
                    buffer[x] = r.getSample(x, y, 0)*1.0/MAX_SAMPLE;
            }
            double[] outBuffer = propagate.get(inBuffer, cache);
            WritableRaster r = out.getRaster();
            int res = cache.outM.getRes();
            for(int x = 0; x < res; ++x) {
                long v = Math.round(outBuffer[x]*outScale*MAX_SAMPLE);
                if(v > MAX_SAMPLE)
                    v = MAX_SAMPLE;
                r.setSample(x, y, 0, v);
            }
        }
        return out;
    }
    public static BufferedImage resize(BufferedImage image, int width, int height) {
        BufferedImage resizedImage = new BufferedImage(width, height,
                IMAGE_TYPE);
        Graphics2D g = resizedImage.createGraphics();
        g.setComposite(AlphaComposite.Src);
        g.setRenderingHint(RenderingHints.KEY_INTERPOLATION,
        RenderingHints.VALUE_INTERPOLATION_BILINEAR);
        g.setRenderingHint(RenderingHints.KEY_RENDERING,
        RenderingHints.VALUE_RENDER_QUALITY);
        g.setRenderingHint(RenderingHints.KEY_ANTIALIASING,
        RenderingHints.VALUE_ANTIALIAS_ON);
        g.drawImage(image, 0, 0, width, height, null);
        g.dispose();
        return resizedImage;
    }
}
